import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

moive = ['Love Movie', 'Action Movie']

T = np.array([[3, 104, 0],
     [2, 100, 0],
     [1, 81, 1,],  # for test, label 0 => 1
     [101, 10, 1],
     [99, 5, 1],
     [98, 2, 1]], dtype=np.float64)
m = len(T)

x = np.array([18, 90])
k = 3

cmap = plt.cm.get_cmap('rainbow', 2)
idx_neg = T[:, 2] == 0
idx_pos = np.invert(idx_neg)
m_neg = np.sum(idx_neg)
m_pos = np.sum(idx_pos)
c00 = cmap(0)
c0 = np.tile(cmap(0), m_neg)
plt.scatter(T[idx_neg, 0], T[idx_neg, 1], c=np.atleast_2d(cmap(0)), label='negative')
plt.scatter(T[idx_pos, 0], T[idx_pos, 1], c=np.atleast_2d(cmap(1)), label='positive')
plt.scatter(x[0], x[1], c='b')
plt.legend()

dis_vector = ((T[:, :2] - x)**2).sum(axis=1)**0.5  # vector (6,)
print(dis_vector.shape)
idx_sorted = dis_vector.argsort()
print(idx_sorted)

label_head = T[idx_sorted[:k], -1]
print(label_head)

labels = pd.Series(label_head)
labels_cnt = labels.value_counts()
h = most_freq_label = int(labels_cnt.index[0])

print(f'Hypothesis = {h} ({moive[h]})')

plt.show()
